#import scipy
from numpy import *
from scipy.integrate import quad, dblquad
from scipy.special import jn_zeros, jnp_zeros, j0, j1, jn, jvp
from scipy.optimize import brentq, fmin, anneal, anderson, broyden1
import matplotlib.pyplot as plt
from scipy.linalg import inv
import DESolver
import time
from pylab import save, load

def j1_zero(n):
    return brentq(j1, pi*(n+0.25-0.05), pi*(n+0.25+0.005))

def j1p_zero(n):
    return brentq(lambda x: jvp(1,x), pi*(n-0.25-0.17), pi*(n-0.25+0.17))

def j1_zeros(ka):
    """ return a list of zeros of the bessel fcn j1(x) up to the value of ka """
    chi = []
    n = 1
    j1z = j1_zero(n)
    while j1z < ka:
        chi.append(j1z)
        n=n+1
        j1z = j1_zero(n)
    return chi

def j1p_zeros(ka):
    """ return a list of zeros of the bessel fcn j1p(x) up to the value of ka """
    chi = []
    n = 1
    j1z = j1p_zero(n)
    while j1z < ka:
        chi.append(j1z)
        n=n+1
        j1z = j1p_zero(n)
    return chi

def j1p_zerosi(N):
    """ return a list of N zeros of the derivative of the bessel fcn j1p(x) """
    chi = []
    for n in range(1,N+1):
        chi.append(j1p_zero(n))
    return chi

def j1_zerosi(N):
    """ return a list of N zeros of the bessel fcn j1p(x) """
    chi = []
    for n in range(1,N+1):
        chi.append(j1_zero(n))
    return chi


def mrad_cyl(freq, theta, phi, a0, b, er=1.0, ur=1.0, sige=0.0, rhom=0.0):
    """ calcuate radiation patterns from on cyl mode coefficients b"""
    
    def Il_f(theta,phi,n):
        # WARNING: propagation constant k must be real (no losses) since bessel function jn does not handle complex numbers
        z = k.real*a0*sin(theta)
        if z==0.0:
            return array([a0*pi*j1(krho_f[n]*a0),0.0])
        else:
            B = pi*j0(z)
            A = 2.0*pi*j1(z)/z-B
            return a0*j1(krho_f[n]*a0)*array([A*cos(2*phi+pi/2.0),A*cos(2.0*phi)+B])
    
    def Is_f(theta,phi,n):
        # WARNING: propagation constant k must be real (no losses) since bessel function jn does not handle complex numbers
        rx = sin(theta)*cos(phi)
        ry = sin(theta)*sin(phi)
        if k.real*sin(theta) == krho_f[n]:
            exp1 = pi*k*sin(phi)*a0**2
            exp2 = j0(krho_f[n]*a0)**2  
            exp3 = j1(krho_f[n]*a0)**2
            exp4 = 2.0/krho_f[n]/a0*j0(krho_f[n]*a0)*j1(krho_f[n]*a0)
            exp5 = exp1*(exp2+exp3-exp4)
            return exp5*array([rx,ry])
        else:    
            exp1 = 2*pi*k*a0/(krho_f[n]**2-(k**2*sin(theta)**2))
            exp2 = krho_f[n]*jn(2,krho_f[n]*a0)*j1(k.real*sin(theta)*a0)
            exp3 = k*sin(theta)*j1(krho_f[n]*a0)*jn(2,k.real*sin(theta)*a0)
            exp4 = exp1*(exp2-exp3)*sin(phi)
            return exp4*array([rx,ry])
    
    def I_f(theta,phi,n):
        return Is_f(theta,phi,n)+Il_f(theta,phi,n)
    
    def I_a(theta,phi,n):
        # WARNING: propagation constant k must be real (no losses) since bessel function jn does not handle complex numbers
        rx = sin(theta)*cos(phi)
        ry = sin(theta)*sin(phi)
        if k.real*sin(theta) == krho_a[n]:
            exp1 = pi*k*sin(phi)*a0**2
            exp2 = j0(krho_a[n]*a0)**2  
            exp3 = j1(krho_a[n]*a0)**2
            exp4 = 2.0/krho_a[n]/a0*j0(krho_a[n]*a0)*j1(krho_a[n]*a0)
            exp5 = exp1*(exp2+exp3-exp4)
            return exp5*array([rx,ry])
        else:    
            exp1 = 2.0*pi*k*a0/(krho_a[n]**2-(k**2*sin(theta)**2))
            exp2 = krho_a[n]*jn(2,krho_a[n]*a0)*j1(k.real*sin(theta)*a0)
            exp3 = k.real*sin(theta)*j1(krho_a[n]*a0)*jn(2,k.real*sin(theta)*a0)
            exp4 = exp1*(exp2-exp3)*cos(phi)
            return exp4*array([rx,ry])
    
    def zxI_a(theta,phi,n):
        i_a = I_a(theta,phi,n)
        return array([-i_a[1],i_a[0]])
    
    def zxI_f(theta,phi,n):
        i_f = I_f(theta,phi,n)
        return array([-i_f[1],i_f[0]])

    def Nt(theta,phi):
        nt_f = 0.0
        for n in range(0,N_f):
            nt_f += -c_f[n]/sqrt(Z_f[n])*zxI_f(theta,phi,n)*b_f[n]
        nt_a = 0.0
        for n in range(0,N_a):
            nt_a += c_a[n]/sqrt(Z_a[n])*I_a(theta,phi,n)*b_a[n]
        return nt_f + nt_a
    
    def Lt(theta,phi):
        lt_f = 0.0
        for n in range(0,N_f):
            lt_f += c_f[n]*sqrt(Z_f[n])*I_f(theta,phi,n)*b_f[n]
        lt_a = 0.0
        for n in range(0,N_a):
            lt_a += c_a[n]*sqrt(Z_a[n])*zxI_a(theta,phi,n)*b_a[n]
        return lt_f + lt_a
    
    def Nsph(theta,phi):
        nt = Nt(theta,phi)
        ntheta = cos(theta)*cos(phi)*nt[0]+cos(theta)*sin(phi)*nt[1]
        nphi = -sin(phi)*nt[0]+cos(phi)*nt[1]
        return array([ntheta,nphi])
    
    def Lsph(theta,phi):
        lt = Lt(theta,phi)
        ltheta = cos(theta)*cos(phi)*lt[0]+cos(theta)*sin(phi)*lt[1]
        lphi = -sin(phi)*lt[0]+cos(phi)*lt[1]
        return array([ltheta,lphi])
    
    def E_ff(theta,phi,r=1.0):
        lsph = Lsph(theta,phi)
        nsph = Nsph(theta,phi)
        etheta_ff = -1j*k*exp(-1j*k*r)/4.0/pi/r*(lsph[1]+nsph[0]*heta)
        ephi_ff =    1j*k*exp(-1j*k*r)/4.0/pi/r*(lsph[0]-nsph[1]*heta)
        return array([etheta_ff,ephi_ff])
    
    def Directivity(theta,phi):
        r = 1.0
        return 2.0*pi/P.real/heta.real*(abs(E_ff(theta,phi,r)[0])**2+abs(E_ff(theta,phi,r)[1])**2)*r**2

    def CrossPol(theta,phi):
        r = 1.0
        return 0.0
    
    D = vectorize(Directivity)
    XPol = vectorize(CrossPol)

    # constants
    c = 299792458.0                 # speed of light (m/s)
    omega = 2.0*pi*freq             # radian frequency (rad/s)
    u0 = 4.0*pi*1e-7                # free space permeability (H/m)
    e0 = 1.0/u0/c/c                 # free space permittivity (F/m)
    mu = u0*ur                      # permeability in media (H/m)
    eps = e0*er                     # permittivity in media (F/m)
    rhom_ = 1.0-1j*rhom/omega/mu    # magnetic loss tangent
    sige_ = 1.0-1j*sige/omega/eps   # electric loss tangent
    mu_ = mu*rhom_                  # complex permeability 
    eps_ = eps*sige_                # complex permittivity
    
    b_f = b[0::2]
    b_a = b[1::2]
    
    k = omega*sqrt(eps_*mu_)        # free space propagation constant (rad/m)
    N_a = len(b_a)
    N_f = len(b_f)
    chi_a = array(j1_zerosi(N_a))
    chi_f = array(j1p_zerosi(N_f))
    
    # calculate propagation constants
    krho_a = chi_a/a0                # propagation constant of TM modes in rho direction (rad/m)
    krho_f = chi_f/a0                # propagation constant of TE modes in rho direciton (rad/m)
    kz_a = -1j*sqrt(krho_a**2-k**2) # propagation constant of TM modes in z direction (rad/m)
    kz_f = -1j*sqrt(krho_f**2-k**2) # propagation constant of TE modes in z direction (rad/m)
        
    #  calculate impedances
    heta = sqrt(mu_/eps_)           # impedance of media
    Z_f = omega*mu_/kz_f            # TM1n mode impedance
    Z_a = kz_a/omega/eps_           # TE1n mode impedance
    ZZ_a = Z_a/abs(Z_a)             # angle of TM impedance
    ZZ_f = Z_f/abs(Z_f)             # angle of TE impedance
        
    # calculate normalization constants
    Irho_a = 0.5*a0**2*(j1(krho_a*a0)**2.0-j0(krho_a*a0)*jn(2,krho_a*a0))
    Irho_f = 0.5*a0**2*(j1(krho_f*a0)**2.0-j0(krho_f*a0)*jn(2,krho_f*a0))
    c_a = sqrt(1.0/Irho_a/pi)/krho_a
    c_f = sqrt(1.0/Irho_f/pi)/krho_f

    # power calculations of radiating fields
    P =  0.5*(sum(ZZ_a*abs(b_a)**2) + sum(ZZ_f*abs(b_f)**2))
    
    return D(theta,phi), XPol(theta,phi)


def gauss2cyl(freq, N, a, gbw=pi/4.0, xdB=10, er=1.0, ur=1.0, sige=0.0, rhom=0.0):
    """ calculate cyl mode coefficients that best represent a fundamental gaussian mode """
    
    def Ig1_f(n):
        return quad(lambda rho: rho*jvp(1, krho_f[n]*rho)*exp(-rho**2/w0**2), 0, a)

    def Ig2_f(n):
        return quad(lambda rho: j1(krho_f[n]*rho)*exp(-rho**2/w0**2), 0, a)

    def Ig1_a(n):
        return quad(lambda rho: rho*jvp(1, krho_a[n]*rho)*exp(-rho**2/w0**2), 0, a)

    def Ig2_a(n):
        return quad(lambda rho: j1(krho_a[n]*rho)*exp(-rho**2/w0**2), 0, a)   

    # constants
    c = 299792458.0                 # speed of light (m/s)
    omega = 2.0*pi*freq             # radian frequency (rad/s)
    u0 = 4.0*pi*1e-7                # free space permeability (H/m)
    e0 = 1.0/u0/c/c                 # free space permittivity (F/m)
    mu = u0*ur                      # permeability in media (H/m)
    eps = e0*er                     # permittivity in media (F/m)
    rhom_ = 1.0-1j*rhom/omega/mu    # magnetic loss tangent
    sige_ = 1.0-1j*sige/omega/eps   # electric loss tangent
    mu_ = mu*rhom_                  # complex permeability 
    eps_ = eps*sige_                # complex permittivity
    
    k = omega*sqrt(eps_*mu_)        # free space propagation constant (rad/m)
    # calculate the number of propagating modes
    if N%2:
        N_f, N_a = N/2+1, N/2
    else:
        N_f, N_a = N/2, N/2
    
    chi_a = array(j1_zerosi(N_a))
    chi_f = array(j1p_zerosi(N_f))
#    chi_f = array(j1p_zerosi(len(chi_a)))
            
#    N_a = len(chi_a)
#    N_f = len(chi_f)
    
    # calculate propagation constants
    krho_a = chi_a/a                # propagation constant of TM modes in rho direction (rad/m)
    krho_f = chi_f/a                # propagation constant of TE modes in rho direciton (rad/m)
    kz_a = -1j*sqrt(krho_a**2-k**2) # propagation constant of TM modes in z direction (rad/m)
    kz_f = -1j*sqrt(krho_f**2-k**2) # propagation constant of TE modes in z direction (rad/m)

#    print "kz_a:", kz_a
#    print "kz_f:", kz_f
    
    #  calculate impedances
    heta = sqrt(mu_/eps_)           # impedance of media
    Z_f = omega*mu_/kz_f            # TM1n mode impedance
    Z_a = kz_a/omega/eps_           # TE1n mode impedance
    
    # calculate normalization constants
    Irho_a = 0.5*a**2*(j1(krho_a*a)**2.0-j0(krho_a*a)*jn(2,krho_a*a))
    Irho_f = 0.5*a**2*(j1(krho_f*a)**2.0-j0(krho_f*a)*jn(2,krho_f*a))
    c_a = sqrt(1.0/Irho_a/pi)/krho_a
    c_f = sqrt(1.0/Irho_f/pi)/krho_f
    
    # calculate the beam waist
    w0 = sqrt(abs(xdB)*log(10.0)/5.0)/k/sin(gbw/2.0);

    # caluclate waveguide mode coefficients  
    b_a = array(zeros(N_a))
    b_f = array(zeros(N_f))
    for n in range(0,N_f):
        b_f[n] = -c_f[n]*sqrt(Z_f[n])*sqrt(2.0*heta/pi)*pi/w0*(krho_f[n]*Ig1_f(n)[0]+Ig2_f(n)[0])
    for n in range(0,N_a):
        b_a[n] = -c_a[n]*sqrt(Z_a[n])*sqrt(2.0*heta/pi)*pi/w0*(krho_a[n]*Ig1_a(n)[0]+Ig2_a(n)[0])
    b = empty(N_a+N_f)
    b[0::2] = b_f
    b[1::2] = b_a
    b = b/sqrt(sum((abs(b))**2))
    return b, w0


def st_c1ce(freq,a1,a2,l1=0.0,l2=0.0,er1=1.0,er2=1.0,ur1=1.0,ur2=1.0,sige1=0.0,sige2=0.0,rhom1=0.0,rhom2=0.0,f_fmax=5.0,ca=0.01):
    """ Models circular waveguide step with different dielectrics in each region
    
    input:
        freq:   freqeuncy of operation
        a1:     radius of waveguide in region 1 (m)
        a2:     radius of waveguide in region 2 (m)
        l1:     length of waveguide in region 1 (m)
        l2:     length of waveguide in region 2 (m)
        er1:    relative permittivty of material in region 1 (unitless)
        er2:    relative permittivty of material in region 2 (unitless)
        ur1:    relative permeability of material in region 1 (unitless)
        ur2:    relative permeability of material in region 2 (unitless)
        sige1:  electric conductivity of material in region 1 (siemens/m)
        sige2:  electric conductivity of material in region 2 (siemens/m)
        rhom1:  magnetic resistance of material in region 1 (ohm/m)
        rhom2:  magnetic resistance of material in region 2 (ohm/m)
        f_fmax: factor times higher than maximum frequency of interest
        ca:  connection accuracy, all modes below this level are discarded
        
    returns the scattering matrix
    """
    def P_FF(i,j):
        """ coupling form TE of region 2 to TE of region 1 """
        k1 = krho_f1[i]
        k2 = krho_f2[j]
        if k1==k2:
            if i==j:
                return sqrt(kz_f1[i]/kz_f2[j]*mu_2/mu_1)
            else:
                return 0.0
        else:
            K = c_f2[j]*c_f1[i]*sqrt(Z_f2[j])/sqrt(Z_f1[i])*a2*pi/((k1/k2)**2-1.0)
            return K*(k2*j0(k2*a2)*j1(k1*a2)-k1*j0(k1*a2)*j1(k2*a2))

    def P_AA(i,j):
        """ coupling form TM of region 2 to TM of region 1 """
        k1 = krho_a1[i]
        k2 = krho_a2[j]
        if k1==k2:
            if i==j:
                return sqrt(kz_a2[j]/kz_a1[i]*eps_1/eps_2)
            else:
                return 0.0
        else:
            K = c_a2[j]*c_a1[i]*sqrt(Z_a2[j])/sqrt(Z_a1[i])*a2*pi*k2
            return K/(1.0-(k2/k1)**2)*j0(k2*a2)*j1(k1*a2) 
    
    def P_FA(i,j):
        """ coupling form TE of region 2 to TM of region 1 """
        if a1==a2:
            return 0.0
        else:
            K = c_f2[j]*c_a1[i]*sqrt(Z_f2[j])/sqrt(Z_a1[i])*pi
            return K*j1(krho_f2[j]*a2)*j1(krho_a1[i]*a2)
    
    def P_AF(i,j):
        """ coupling form TM of region 2 to TE of region 1 """
        return 0.0

    # interchange region 1 and region 2 dimensions and properties
    reverse = False
    if a1 < a2:
        reverse = True
        a1,a2 = a2,a1
        l1,l2 = l2,l1
        er1,er2 = er2,er1
        ur1,ur2 = ur2,ur1
        sige1,sige2 = sige2,sige1
        rhom1,rhom2 = rhom2,rhom1        
    
    # constants
    c = 299792458.0                 # speed of light (m/s)
    omega = 2.0*pi*freq             # radian frequency (rad/s)
    
    u0 = 4.0*pi*1e-7                # free space permeability (H/m)
    e0 = 1.0/u0/c/c                 # free space permittivity (F/m)
    
    mu1 = u0*ur1                      # permeability in media (H/m)
    eps1 = e0*er1                     # permittivity in media (F/m)
    rhom_1 = 1.0-1j*rhom1/omega/mu1    # magnetic loss tangent
    sige_1 = 1.0-1j*sige1/omega/eps1  # electric loss tangent
    mu_1 = mu1*rhom_1                  # complex permeability 
    eps_1 = eps1*sige_1                # complex permittivity
    
    mu2 = u0*ur2                      # permeability in media (H/m)
    eps2 = e0*er2                     # permittivity in media (F/m)
    rhom_2 = 1.0-1j*rhom2/omega/mu2    # magnetic loss tangent
    sige_2 = 1.0-1j*sige2/omega/eps2   # electric loss tangent
    mu_2 = mu2*rhom_2                  # complex permeability 
    eps_2 = eps2*sige_2                # complex permittivity
    
    k1 = omega*sqrt(eps_1*mu_1)        # free space propagation constant (rad/m)
    k2 = omega*sqrt(eps_2*mu_2)        # free space propagation constant (rad/m)

    # determine the total number of modes in each region
    ka1 = f_fmax*k1.real*a1
    ka2 = f_fmax*k2.real*a2
    
    chi_a1 = array(j1_zeros(ka1))
    chi_a2 = array(j1_zeros(ka2))
    chi_f1 = array(j1p_zerosi(len(chi_a1)))
    chi_f2 = array(j1p_zerosi(len(chi_a2)))
        
    # calculate the number of TE/TM modes
    N1 = len(chi_a1)
    N2 = len(chi_a2)

    # calculate the propagation constants
    krho_a1 = chi_a1/a1                # propagation constant of TM modes in rho direction (rad/m)
    krho_f1 = chi_f1/a1                # propagation constant of TE modes in rho direciton (rad/m)
    kz_a1 = -1j*sqrt(krho_a1**2-k1**2) # propagation constant of TM modes in z direction (rad/m)
    kz_f1 = -1j*sqrt(krho_f1**2-k1**2) # propagation constant of TE modes in z direction (rad/m)
    
    krho_a2 = chi_a2/a2                # propagation constant of TM modes in rho direction (rad/m)
    krho_f2 = chi_f2/a2                # propagation constant of TE modes in rho direciton (rad/m)
    kz_a2 = -1j*sqrt(krho_a2**2-k2**2) # propagation constant of TM modes in z direction (rad/m)
    kz_f2 = -1j*sqrt(krho_f2**2-k2**2) # propagation constant of TE modes in z direction (rad/m)
    
    #  calculate impedances
    Z_f1 = omega*mu_1/kz_f1            # TM1n mode impedance
    Z_a1 = kz_a1/omega/eps_1           # TE1n mode impedance
    
    Z_f2 = omega*mu_2/kz_f2            # TM1n mode impedance
    Z_a2 = kz_a2/omega/eps_2           # TE1n mode impedance
    
    # calculate normalization constants
    c_a1 = 1.0/sqrt(0.5*(-j0(chi_a1)*jn(2,chi_a1))*pi)/chi_a1
    c_f1 = 1.0/sqrt(0.5*(j1(chi_f1)**2.0-j0(chi_f1)*jn(2,chi_f1))*pi)/chi_f1
    
    c_a2 = 1.0/sqrt(0.5*(-j0(chi_a2)*jn(2,chi_a2))*pi)/chi_a2
    c_f2 = 1.0/sqrt(0.5*(j1(chi_f2)**2.0-j0(chi_f2)*jn(2,chi_f2))*pi)/chi_f2

    # order the modes according to the inequality in Abramowitz and Stegun 9.5.2
    # m <= j'm,1 < jm,1 < j'm,2 < jm,2 < j'm,3 ...
    # for our purposes m=1
    # 1 <= j'1,1 < j1,1 < j'1,2 < j1,2 < j'1,3 ... since 1<j'1,1 = 1.841
    # thereore the TE and TM modes interleave.  Since matrices start with index 0
    # TE indices are even i = 0,2,4 ...
    # TM indices are odd j = 1,3,5 ...
    kz1 = empty(2*N1, dtype=complex)
    kz1[0::2] = kz_f1
    kz1[1::2] = kz_a1

    kz2 = empty(2*N2, dtype=complex)
    kz2[0::2]=kz_f2
    kz2[1::2]=kz_a2

    P = zeros((2*N1, 2*N2), complex)
    for i in arange(0,N1):
        for j in arange(0,N2):
            m = 2*i
            n = 2*j
            P[m,n] = P_FF(i,j)
            P[m,n+1] = P_AF(i,j)
            P[m+1,n] = P_FA(i,j)
            P[m+1,n+1] = P_AA(i,j)

    Q = P.T
    # calculate scattering matrix partitions orderend in column major form
    U1 = eye(2*N1)
    U2 = eye(2*N2)
    W = inv(U1 + dot(P,Q))
    V = inv(U2 + dot(Q,P))
    S11 = dot(W,dot(P,Q)-U1)
    S12 = dot(W,2*P)
    S21 = dot(V,2*Q)
    S22 = dot(V,U2-dot(Q,P))

    # calculate the length extensions of s-parameters
    D1 = exp(-1j*kz1*l1)
    D2 = exp(-1j*kz2*l2)

    # array on indices of propagating modes, used for modal radiation
    ip1 =  nonzero(kz1.real>0.0)[0]
    ip2 =  nonzero(kz2.real>0.0)[0]

    # array of indices of connected modes
    ic1 = nonzero(abs(D1)>=ca)[0]
    ic2 = nonzero(abs(D2)>=ca)[0]

    # resize s-matrix to only include modes that contribute
    D1 = D1[ix_(ic1)]
    D2 = D2[ix_(ic2)]
    S11 = S11[ix_(ic1,ic1)]
    S12 = S12[ix_(ic1,ic2)]
    S21 = S21[ix_(ic2,ic1)]
    S22 = S22[ix_(ic2,ic2)]
    
    # add length extensions to s-matrix
    S11 = (D1*((D1*S11).T)).T
    S12 = (D1*((D2*S12).T)).T
    S21 = (D2*((D1*S21).T)).T
    S22 = (D2*((D2*S22).T)).T
  
    # return value of matrix is column major order
    # S11, S21, S12, S22
    #[[ S11, S12]
    # [ S21, S22]]
    if reverse==True: #a1<a2
        return S22, S12, S21, S11, ip2, ip1
    else: # a1>=a2
        return S11, S21, S12, S22, ip1, ip2
    
def connect2p(S11A, S21A, S12A, S22A, S11B, S21B, S12B, S22B):
    """ Calculate the connection of two ports forming a new scattering matrix"""
    
    U = identity(shape(S11B)[0])
    V = dot(S12A, inv(U - dot(S11B, S22A)))
    S11 = S11A + dot(dot(V, S11B), S21A)
    S12 = dot(V, S12B)
    W = dot(S21B, inv(U - dot(S22A, S11B)))
    S21 = dot(W, S21A)
    S22 = S22B + dot(dot(W, S22A), S12B)
    return (S11, S21, S12, S22)

def Directivity_gaussian(theta, k, w0):
    Umax = sqrt(2)*abs(k)*w0
    return Umax*0.5*exp(-(w0*k*sin(theta)/2.0)**2)*(1.0+cos(theta))  

def plot_geometry(r,l):
    # used for plotting the structure
    rn = zeros(r.shape[0]*2)
    rn[0::2] = r
    rn[1::2] = r
    rn *= 1000.0 # convert to mm
    zn = zeros(r.shape[0]*2)
    zn[1::2] = cumsum(l)
    zn[2:-1:2] = cumsum(l[:-1])
    zn[0] += l[0]/2.0
    zn[-1] -= l[-1]/2.0
    zn *= 1000.0 # convert to mm
    
    plt.plot(zn,rn,'b')
    plt.plot(zn,-rn,'b')
    plt.plot([zn[-1],zn[-1]],[rn[-1],-rn[-1]],'b',[zn[0],zn[0]],[rn[0],-rn[0]],'b')
    plt.axis('equal')
    plt.xlabel('z axis (mm)')
    plt.ylabel(r'$\rho$ axis (mm)')
    plt.title('Geometry of radiating structure')

def func(x0, *args):
    # gaussian beam properties
    gbw0 = 47.0/180.0*pi
    xdB0 = 10
    
    # convergence/accuracy parameters
    f_fmax = 5.0
    ca = 0.1
    
    freq = args[0]
    r = args[1]
    r_indx = args[2]
    l = args[3]
    l_indx = args[4]
    display = args[5]

    # no. of regions
    Nr = len(r)
    r[ix_(r_indx)]=x0[:len(r[ix_(r_indx)])]
    l[ix_(l_indx)]=x0[len(r[ix_(r_indx)]):]
       
    er = 1.0
    ur = 1.0
    
    geom = zeros((6,Nr))
    geom[0] = r
    geom[1] = l
    geom[2] = ones([1,Nr])*er # relative permitivitties
    geom[3] = ones([1,Nr])*ur # relative permeabilities
    geom[4] = zeros([1,Nr]) # sige
    geom[5] = zeros([1,Nr]) # rhom
    
    
    # variables to store s-parameters as a function of frequency
    s11 = []
    s21 = []
    s22 = []
    directivity = []
    
    # gaussian targets
    # list of varying size arrays of mode coefficients at a particular frequency
    b_frq = [] 
#    b_gauss = []
#    w0_frq = []

    m=0
    for f in freq:
        # initial step
        n=0
        (S11_0, S21_0, S12_0, S22_0, ip1, ip2) = st_c1ce(f, 
           a1=geom[0,n], a2=geom[0,n+1], 
           l1=geom[1,n]/2.0, l2=geom[1,n+1]/2.0, 
           er1=geom[2,n], er2=geom[2,n+1], 
           ur1=geom[3,n], ur2=geom[3,n+1], 
           sige1=geom[4,n], sige2=geom[4,n+1],
           rhom1=geom[5,n], rhom2=geom[5,n+1],
           f_fmax=f_fmax, ca=ca)
        # iterate over the remaining steps
        for n in range(1,Nr-1):
            (S11_1, S21_1, S12_1, S22_1, ip1, ip2) = st_c1ce(f, 
               a1=geom[0,n], a2=geom[0,n+1], 
               l1=geom[1,n]/2.0, l2=geom[1,n+1]/2.0, 
               er1=geom[2,n], er2=geom[2,n+1], 
               ur1=geom[3,n], ur2=geom[3,n+1], 
               sige1=geom[4,n], sige2=geom[4,n+1],
               rhom1=geom[5,n], rhom2=geom[5,n+1],
               f_fmax=f_fmax, ca=ca)
            (S11_0, S21_0, S12_0, S22_0) = connect2p(S11_0, S21_0, S12_0, S22_0, S11_1, S21_1, S12_1, S22_1)    
        b_frq.append(S21_0[ix_(ip2),0][0])
#        b_gauss.append(gauss2cyl(freq[m], len(ip2), geom[0,-1], gbw=gbw0, xdB=xdB0, er=geom[2,-1], ur=geom[3,-1], sige=geom[4,-1], rhom=geom[5,-1])[0])
        directivity.append(mrad_cyl(freq[m], 0.0, 0.0, geom[0,-1], b_frq[m])[0])
        
        s11.append(S11_0[0,0])
        s21.append(S21_0[0,0])
        s22.append(S22_0[0,0])
        #print "freq:", freq[m]*1e-9, "GHz"
        m=m+1

    s11 = array(s11)
    s21 = array(s21)
    s22 = array(s22)
    
    # use this error function when optimizing with gain
    error = []
    directivity = array(directivity)
    q = 1.0-abs(s11)**2 # mismatch efficiency
    gain = directivity*q
    error = max(1.0/gain)
    print 'error from gain:', error 
    
    # use this error function when optimizing with gaussian targets
    # iterate over list of arrys
#    error = []
#    for n in range(len(freq)):
#        error.append(sqrt(sum((abs(b_frq[n])*abs(b_frq[n]-abs(b_gauss[n])))**2)))         
#    error = array(error)
#    error = 1.0/sqrt(sum(error*error))
#    print 'error from gaussian:', error

    # save optimization results to file
    save("opt/horn_opt_simplex_radius.txt", r)
    save("opt/horn_opt_simplex_length.txt", l)
    print "worst case gain:", 10*log10(min(gain))


    if display==True:
        #===============================================================================
        # plot geometry of structure
        #===============================================================================
        print 'plotting structure geometry'
        plt.subplot(2,2,1)
        plot_geometry(geom[0], geom[1])
        
        #===============================================================================
        # plot scattering parameters
        #===============================================================================
        print 'plotting scattering parameters'
        plt.subplot(2,2,2)
        s11 = array(s11)
        s21 = array(s21)
        s22 = array(s22)
        s11 = 20.0*log10(abs(s11))
        s21 = 20.0*log10(abs(s21))
        s22 = 20.0*log10(abs(s22))    
        plt.plot(freq/1e9,s11,'r',freq/1e9,s21,'b')
        plt.grid()
        plt.axis([14.4,15.35,-40, 0])
        plt.xlabel('freqeuncy (GHz)')
        plt.ylabel('S-parameters (dB)')
        plt.title('Scattering Parameters')
        plt.legend(['S11','S21'])
        
        #===============================================================================
        # #plot radiation patterns
        #===============================================================================
        print 'calculating radiation patterns'
        plt.subplot(2,2,3)
        theta = linspace(0,pi,180/2)
        for n in range(len(freq)):
            (D_phi0, xpol_phi0) = mrad_cyl(freq[n], theta, 0.0, geom[0,-1], b_frq[n])
            (D_phi90, xpol_phi90) = mrad_cyl(freq[n], theta, pi/2.0, geom[0,-1], b_frq[n])
#            (D_phi45, xpol_phi45) = mrad_cyl(freq[n], theta, pi/4.0, geom[0,-1], b_frq[n])
            D_phi0 *= q[n]
            D_phi90 *= q[n]
            plt.plot(theta*180/pi, 10.0*log10(D_phi0),'r', theta*180/pi, 10.0*log10(D_phi90),'b')
#            # gaussian targets
#            (Dg_phi0, xpol_phi0) = mrad_cyl(freq[n], theta, 0.0, geom[0,-1], b_gauss[n])
#            (Dg_phi90, xpol_phi90) = mrad_cyl(freq[n], theta, pi/2.0, geom[0,-1], b_gauss[n])
#            plt.plot(theta*180/pi, 10.0*log10(Dg_phi0),'r--', theta*180/pi, 10.0*log10(Dg_phi90),'b--')
#            print 'freq:', freq[n]*1e-9, 'GHz'
            
        plt.grid()
        plt.xlabel(r'$\theta$ (deg)')
        plt.ylabel('realized gain (dBi)')
        plt.title(r'Realized gain($\theta$) elevation cuts')
        plt.legend([r'$\phi=0$',r'$\phi=\frac{\pi}{2}$'])
        plt.axis([0, 90, -40, 25])
        
        
        #===============================================================================
        # plot max directivity (theta = 0, phi = 0) vs frequency
        #===============================================================================
        print 'plotting maximum gain vs frequency'
        plt.subplot(2,2,4)
        plt.plot(freq/1e9,10*log10(gain))
        plt.grid()
        #plt.axis([14.4,15.35,-30, 0])
        plt.xlabel('freqeuncy (GHz)')
        plt.ylabel('realized gain (dBi)')
        plt.title(r'Maximum realized gain($\theta=0$, $\phi=0$)')
        plt.show()

    return error

if __name__ == '__main__':




#    # radius of end points of horn
    r0 = 7.4198e-3          # input radius
    rn = 43.18e-3           # output radius

    c = 299792458
#    fmin = 14.4e9
#    fmax = 15.35e9
    f0 = 14.875e9/2
    lam = c/f0
    lam_c = 2.0*pi*r0/1.84118
    kc = 2.0*pi/lam_c 
    
#    r0 = 1.5*lam_c/pi
    L = 4*lam            # length
    
    d = 0.5                 # slot pitch-to-width ratio
    p = lam_c/10.0            # slot pitch
    N = floor(L/p)
    w = d*p                 # slot width
    t = p-w                 # tooth width
    
    # radius steps of horn
    r  = linspace(r0,rn,N+1)
    # add corrugations
#    kk = exp(1.0/2.114/(kc*rn)**1.134)
#    r[1::2] = linspace(kk*lam_c/4.0, kk*lam_c/4.0, len(r[1::2])) + r[1::2] 
    r[1:5] = 10.0e-3
#    r = append(ones(1)*r0,r) # add 5 regions onto front

    l = ones(len(r))
    l[1::2] *= w
    l[2::2] *= t      
        
#    r = load("opt/horn_opt_simplex_radius_5.txt")
#    l = load("opt/horn_opt_simplex_length_5.txt")
    l[0] = 25.0e-3
    l[-1] = L/len(r)*2
    r_indx = arange(0,len(r)-1)
    l_indx = arange(1,len(l)-1)
    x0 = append(r[ix_(r_indx)],l[ix_(l_indx)])

    print 'total length:', sum(l)*25.4    
#    plot_geometry(r,l)
#    plt.show()
#    freq = linspace(14.4e9,15.35e9,3)
#    print func(x0, *(freq,r,r_indx,l,l_indx,True))
    freq = linspace(14.4e9,15.35e9,3)
    xopt = fmin(func, x0, args=(freq,r,r_indx,l,l_indx,False), xtol = 0.00001, ftol=0.00001)
    
